from atm_system import ATMSystem


def display_main_menu():
    """显示主菜单"""
    print("\n" + "*" * 41)
    print("*" + " " * 11 + "欢迎使用银行ATM系统" + " " * 11 + "*")
    print("*" * 41)
    print("*" + " " * 15 + "1. 开户" + " " * 17 + "*")
    print("*" + " " * 15 + "2. 存钱" + " " * 17 + "*")
    print("*" + " " * 15 + "3. 取钱" + " " * 17 + "*")
    print("*" + " " * 15 + "4. 转账" + " " * 17 + "*")
    print("*" + " " * 15 + "5. 查询" + " " * 17 + "*")
    print("*" + " " * 15 + "6. 退出" + " " * 17 + "*")
    print("*" * 41)


def get_user_input(prompt, input_type=str):
    """获取用户输入并进行类型转换"""
    while True:
        try:
            user_input = input(prompt)
            if input_type == int:
                return int(user_input)
            elif input_type == float:
                return float(user_input)
            else:
                return user_input
        except ValueError:
            print("输入格式错误，请重新输入！")


def open_account(atm):
    """开户功能"""
    print("\n--- 开户 ---")

    # 获取账户类型
    print("请选择账户类型：")
    print("1. 普通账户")
    print("2. VIP账户")
    account_type_choice = get_user_input("请选择(1-2): ", int)
    account_type = "普通账户" if account_type_choice == 1 else "VIP账户"

    # 获取用户信息
    name = get_user_input("请输入姓名: ")
    password = get_user_input("请输入密码: ", int)

    print("请输入地址信息：")
    country = get_user_input("国家: ")
    province = get_user_input("省份: ")
    city = get_user_input("城市: ")
    state = get_user_input("街道: ")
    number = get_user_input("门牌号: ")

    # 调用开户方法
    result = atm.create_account(account_type, name, password, country, province, city, state, number)


def deposit_money(atm):
    """存钱功能"""
    print("\n--- 存钱 ---")
    account = get_user_input("请输入账号: ", int)
    amount = get_user_input("请输入存款金额: ", float)

    atm.deposit(account, amount)


def withdraw_money(atm):
    """取钱功能"""
    print("\n--- 取钱 ---")
    account = get_user_input("请输入账号: ", int)
    password = get_user_input("请输入密码: ", int)
    amount = get_user_input("请输入取款金额: ", float)

    result = atm.withdraw(account, password, amount)


def transfer_money(atm):
    """转账功能"""
    print("\n--- 转账 ---")
    from_account = get_user_input("请输入转出账号: ", int)
    from_password = get_user_input("请输入转出账号密码: ", int)
    to_account = get_user_input("请输入转入账号: ", int)
    amount = get_user_input("请输入转账金额: ", float)

    result = atm.transfer(from_account, from_password, to_account, amount)


def query_account_info(atm):
    """查询账户信息"""
    print("\n--- 查询账户信息 ---")
    account = get_user_input("请输入账号: ", int)
    password = get_user_input("请输入密码: ", int)

    atm.get_account_info(account, password)


def main():
    """主程序"""
    atm = ATMSystem()

    while True:
        display_main_menu()
        choice = get_user_input("请选择操作(1-6): ", int)

        if choice == 1:
            open_account(atm)
        elif choice == 2:
            deposit_money(atm)
        elif choice == 3:
            withdraw_money(atm)
        elif choice == 4:
            transfer_money(atm)
        elif choice == 5:
            query_account_info(atm)
        elif choice == 6:
            print("感谢使用银行ATM系统，再见！")
            break
        else:
            print("无效选择，请重新输入！")

        input("\n按回车键继续...")


if __name__ == "__main__":
    main()
